import tkinter as tk
from tkinter import ttk, messagebox, filedialog
import akshare as ak
import pandas as pd
from datetime import datetime
import os
import sys
import matplotlib.pyplot as plt
import threading

# 添加项目根目录到Python路径
sys.path.append(os.path.join(os.path.dirname(__file__), ".."))

# 尝试导入模型相关模块
try:
    from model import Kronos, KronosTokenizer, KronosPredictor
    MODEL_AVAILABLE = True
except ImportError:
    MODEL_AVAILABLE = False
    print("Warning: Model modules not available. Prediction feature will be disabled.")

def plot_prediction(kline_df, pred_df):
    """绘制预测结果"""
    pred_df.index = kline_df.index[-pred_df.shape[0]:]
    sr_close = kline_df['close']
    sr_pred_close = pred_df['close']
    sr_close.name = 'Ground Truth'
    sr_pred_close.name = "Prediction"

    sr_volume = kline_df['volume']
    sr_pred_volume = pred_df['volume']
    sr_volume.name = 'Ground Truth'
    sr_pred_volume.name = "Prediction"

    close_df = pd.concat([sr_close, sr_pred_close], axis=1)
    volume_df = pd.concat([sr_volume, sr_pred_volume], axis=1)

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8), sharex=True)

    ax1.plot(close_df['Ground Truth'], label='Ground Truth', color='blue', linewidth=1.5)
    ax1.plot(close_df['Prediction'], label='Prediction', color='red', linewidth=1.5)
    ax1.set_ylabel('Close Price', fontsize=14)
    ax1.legend(loc='lower left', fontsize=12)
    ax1.grid(True)
    ax1.set_title('Price Prediction', fontsize=16)

    ax2.plot(volume_df['Ground Truth'], label='Ground Truth', color='blue', linewidth=1.5)
    ax2.plot(volume_df['Prediction'], label='Prediction', color='red', linewidth=1.5)
    ax2.set_ylabel('Volume', fontsize=14)
    ax2.set_xlabel('Time', fontsize=14)
    ax2.legend(loc='upper left', fontsize=12)
    ax2.grid(True)
    ax2.set_title('Volume Prediction', fontsize=16)

    plt.tight_layout()
    plt.show()

class StockExtractorApp:
    def __init__(self, root):
        self.root = root
        self.root.title("Stock Data Prediction")
        self.root.geometry("500x480")
        self.root.resizable(False, False)  # 禁止调整窗口大小
        
        # 将窗口居中显示
        self.center_window()
        
        # 创建界面元素
        self.create_widgets()
        
    def center_window(self):
        """将窗口居中显示在屏幕上"""
        self.root.update_idletasks()
        width = self.root.winfo_width()
        height = self.root.winfo_height()
        x = (self.root.winfo_screenwidth() // 2) - (width // 2)
        y = (self.root.winfo_screenheight() // 2) - (height // 2)
        self.root.geometry(f"{width}x{height}+{x}+{y}")
        
    def create_widgets(self):
        # 主框架
        main_frame = ttk.Frame(self.root, padding="20")
        main_frame.pack(fill=tk.BOTH, expand=True)
        
        # 标题
        title_label = ttk.Label(main_frame, text="Stock Data Prediction", font=("Arial", 14, "bold"))
        title_label.pack(pady=(0, 15))
        
        # 输入区域框架
        input_frame = ttk.LabelFrame(main_frame, text="Data Parameters", padding="10")
        input_frame.pack(fill=tk.X, pady=(0, 15))
        
        # 股票代码输入
        ttk.Label(input_frame, text="Stock Code:").grid(row=0, column=0, sticky=tk.W, pady=5)
        self.code_entry = ttk.Entry(input_frame, width=20)
        self.code_entry.grid(row=0, column=1, padx=(10, 0), pady=5, sticky=tk.EW)
        ttk.Label(input_frame, text="e.g., 000001", font=("Arial", 8)).grid(row=1, column=1, padx=(10, 0), sticky=tk.W)
        
        # 开始日期输入
        ttk.Label(input_frame, text="Start Date:").grid(row=2, column=0, sticky=tk.W, pady=(10, 5))
        self.start_entry = ttk.Entry(input_frame, width=20)
        self.start_entry.grid(row=2, column=1, padx=(10, 0), pady=(10, 5), sticky=tk.EW)
        ttk.Label(input_frame, text="YYYY-MM-DD", font=("Arial", 8)).grid(row=3, column=1, padx=(10, 0), sticky=tk.W)
        
        # 结束日期输入
        ttk.Label(input_frame, text="End Date:").grid(row=4, column=0, sticky=tk.W, pady=(10, 5))
        self.end_entry = ttk.Entry(input_frame, width=20)
        self.end_entry.grid(row=4, column=1, padx=(10, 0), pady=(10, 5), sticky=tk.EW)
        ttk.Label(input_frame, text="YYYY-MM-DD", font=("Arial", 8)).grid(row=5, column=1, padx=(10, 0), sticky=tk.W)
        
        # 配置列权重
        input_frame.columnconfigure(1, weight=1)
        
        # 按钮区域
        button_frame = ttk.Frame(main_frame)
        button_frame.pack(fill=tk.X, pady=(10, 15))
        
        # 创建内部按钮框架并居中
        buttons_container = ttk.Frame(button_frame)
        buttons_container.pack(expand=True)
        
        # 提取数据按钮
        self.extract_button = ttk.Button(buttons_container, text="Collect Data", command=self.extract_data)
        self.extract_button.pack(side=tk.LEFT, padx=(0, 10))
        
        # 选择文件按钮
        self.select_file_button = ttk.Button(buttons_container, text="Select File", command=self.select_file)
        self.select_file_button.pack(side=tk.LEFT, padx=(0, 10))
        
        # 清空按钮
        self.clear_button = ttk.Button(buttons_container, text="Clear", command=self.clear_selection)
        self.clear_button.pack(side=tk.LEFT, padx=(0, 10))
        
        # 预测按钮（仅在模型可用时启用）
        self.predict_button = ttk.Button(buttons_container, text="Predict", command=self.perform_prediction, state=tk.DISABLED)
        self.predict_button.pack(side=tk.LEFT, padx=(0, 10))
        
        # 退出按钮
        ttk.Button(buttons_container, text="Exit", command=self.confirm_exit).pack(side=tk.LEFT)
        
        # 状态标签
        self.status_label = ttk.Label(main_frame, text="Ready")
        self.status_label.pack(pady=(0, 5))
        
        # 保存文件路径
        self.saved_file_path = None
        
    def clear_selection(self):
        """清空文件选择并置灰预测按钮"""
        self.saved_file_path = None
        self.predict_button.config(state=tk.DISABLED)
        self.status_label.config(text="Ready - File selection cleared")
        
    def confirm_exit(self):
        """确认退出应用"""
        if messagebox.askokcancel("Confirm Exit", "Are you sure you want to exit the application?"):
            self.root.quit()
        
    def select_file(self):
        """选择CSV文件进行预测"""
        file_path = filedialog.askopenfilename(
            title="Select CSV File",
            filetypes=[("CSV Files", "*.csv"), ("All Files", "*.*")]
        )
        
        if file_path:
            # 验证文件是否存在且是CSV文件
            if os.path.exists(file_path) and file_path.endswith('.csv'):
                self.saved_file_path = file_path
                self.status_label.config(text=f"Selected file: {os.path.basename(file_path)}")
                
                # 如果模型可用，启用预测按钮
                if MODEL_AVAILABLE:
                    self.predict_button.config(state=tk.NORMAL)
                messagebox.showinfo("Success", f"File selected: {os.path.basename(file_path)}")
            else:
                messagebox.showerror("Error", "Please select a valid CSV file")
                self.status_label.config(text="Ready")
    
    def extract_data(self):
        """Collect Data"""
        stock_code = self.code_entry.get().strip()
        start_date = self.start_entry.get().strip()
        end_date = self.end_entry.get().strip()
        
        # Validate inputs
        if not all([stock_code, start_date, end_date]):
            messagebox.showerror("Error", "Please fill in all fields")
            return
            
        # Validate date format
        try:
            datetime.strptime(start_date, "%Y-%m-%d")
            datetime.strptime(end_date, "%Y-%m-%d")
        except ValueError:
            messagebox.showerror("Error", "Incorrect date format, please use YYYY-MM-DD format")
            return
            
        # Update status
        self.status_label.config(text="Collecting data...")
        self.extract_button.config(state=tk.DISABLED)  # 禁用按钮防止重复点击
        self.root.update()
        
        try:
            # Format dates
            start_date_fmt = datetime.strptime(start_date, "%Y-%m-%d").strftime("%Y%m%d")
            end_date_fmt = datetime.strptime(end_date, "%Y-%m-%d").strftime("%Y%m%d")
            
            # Get data
            df = ak.stock_zh_a_hist(symbol=stock_code, period="daily", 
                                  start_date=start_date_fmt, end_date=end_date_fmt)
            
            if df.empty:
                messagebox.showwarning("Warning", "No data retrieved")
                self.status_label.config(text="Ready")
                self.extract_button.config(state=tk.NORMAL)
                return
                
            # Select required columns
            columns_mapping = {
                '日期': 'timestamps',
                '开盘': 'open',
                '最高': 'high',
                '最低': 'low',
                '收盘': 'close',
                '成交量': 'volume',
                '成交额': 'amount'
            }
            
            # Check if columns exist
            available_columns = [col for col in columns_mapping.keys() if col in df.columns]
            if not available_columns:
                messagebox.showerror("Error", "Data columns do not match")
                self.status_label.config(text="Ready")
                self.extract_button.config(state=tk.NORMAL)
                return
                
            # Rename columns
            new_df = df[available_columns].copy()
            new_df.rename(columns={k: v for k, v in columns_mapping.items() if k in available_columns}, inplace=True)
            
            # Save file
            filename = f"XSHG_daily_{stock_code}.csv"
            filepath = os.path.join(os.path.dirname(__file__), filename)
            new_df.to_csv(filepath, index=False)
            
            # 保存文件路径以便后续预测使用
            self.saved_file_path = filepath
            
            # 更新状态
            self.status_label.config(text=f"Success! Data saved to {filename}")
            messagebox.showinfo("Success", f"Successfully exported {len(new_df)} records")
            
            # 如果模型可用，启用预测按钮
            if MODEL_AVAILABLE:
                self.predict_button.config(state=tk.NORMAL)
            
        except Exception as e:
            messagebox.showerror("Error", f"Error collecting data: {str(e)}")
            self.status_label.config(text="Ready")
        finally:
            # 重新启用按钮
            self.extract_button.config(state=tk.NORMAL)
    
    def _run_prediction(self):
        """在后台线程中运行预测功能"""
        try:
            # 加载模型和分词器
            tokenizer = KronosTokenizer.from_pretrained("NeoQuasar/Kronos-Tokenizer-base")
            model = Kronos.from_pretrained("NeoQuasar/Kronos-small")
            
            # 创建预测器
            predictor = KronosPredictor(model, tokenizer, device="cpu", max_context=512)
            
            # 准备数据
            df = pd.read_csv(self.saved_file_path)
            df['timestamps'] = pd.to_datetime(df['timestamps'])
            
            # 设置预测参数
            lookback = min(400, len(df))  # 确保不超过数据长度
            pred_len = min(120, len(df) - lookback)  # 确保有足够的数据进行预测
            
            if pred_len <= 0:
                # 在主线程中显示错误消息
                self.root.after(0, lambda: messagebox.showerror("Error", "Not enough data for prediction. Please extract more data."))
                self.root.after(0, lambda: self.status_label.config(text="Ready"))
                self.root.after(0, lambda: self.predict_button.config(state=tk.NORMAL))
                return
            
            x_df = df.loc[:lookback-1, ['open', 'high', 'low', 'close', 'volume', 'amount']]
            x_timestamp = df.loc[:lookback-1, 'timestamps']
            y_timestamp = df.loc[lookback:lookback+pred_len-1, 'timestamps']
            
            # 执行预测
            pred_df = predictor.predict(
                df=x_df,
                x_timestamp=x_timestamp,
                y_timestamp=y_timestamp,
                pred_len=pred_len,
                T=1.0,
                top_p=0.9,
                sample_count=1,
                verbose=True
            )
            
            # 绘制结果
            kline_df = df.loc[:lookback+pred_len-1]
            
            # 在主线程中更新UI和显示图表
            self.root.after(0, lambda: self.status_label.config(text="Prediction completed"))
            self.root.after(0, lambda: messagebox.showinfo("Success", "Prediction completed and chart displayed"))
            self.root.after(0, lambda: self.predict_button.config(state=tk.NORMAL))
            self.root.after(0, lambda: plot_prediction(kline_df, pred_df))
            
        except Exception as e:
            # 在主线程中显示错误消息
            self.root.after(0, lambda: messagebox.showerror("Error", f"Error performing prediction: {str(e)}"))
            self.root.after(0, lambda: self.status_label.config(text="Ready"))
            self.root.after(0, lambda: self.predict_button.config(state=tk.NORMAL))
    
    def perform_prediction(self):
        """执行预测功能（异步方式）"""
        if not MODEL_AVAILABLE:
            messagebox.showerror("Error", "Prediction model not available")
            return
            
        if not self.saved_file_path or not os.path.exists(self.saved_file_path):
            messagebox.showerror("Error", "No data file available for prediction")
            return
            
        # 更新状态
        self.status_label.config(text="Performing prediction...")
        self.predict_button.config(state=tk.DISABLED)
        
        # 在后台线程中运行预测
        prediction_thread = threading.Thread(target=self._run_prediction)
        prediction_thread.daemon = True  # 设置为守护线程，确保主程序退出时线程也会退出
        prediction_thread.start()

def main():
    root = tk.Tk()
    app = StockExtractorApp(root)
    root.mainloop()

if __name__ == "__main__":
    main()